
import sys
import torch
import torch.utils._pytree as _pytree

print("Testing PyTree patch...")

# 1. Define the patch
_orig_register = _pytree._register_pytree_node
def _compat_register(cls, flatten_fn, unflatten_fn, **_):
    return _orig_register(cls, flatten_fn, unflatten_fn)

# 2. Apply the patch
_pytree.register_pytree_node = _compat_register
print("Patch applied.")

# 3. Import torchvision (which imports transformers -> uses register_pytree_node)
try:
    import torchvision.transforms as transforms
    print("Successfully imported torchvision.transforms")
except AttributeError as e:
    print(f"Failed to import torchvision: {e}")
    sys.exit(1)
except Exception as e:
    print(f"An unexpected error occurred: {e}")
    sys.exit(1)

print("Verification successful!")
